import argparse
import re
import os

script_path = os.path.dirname(os.path.abspath(__file__))

class Inventory:
    def __init__(self, inventory_template_file, output_inventory_file):
        self.inventory_template_file = inventory_template_file
        self.output_inventory_file = output_inventory_file
        with open(os.path.join(script_path, "inventory_template", self.inventory_template_file), 'r',
                  encoding='utf-8') as f:
            self.inventory_lines = f.readlines()
        self.node_pattern = r'(\{node\d+\})'
        self.node_number = self.get_node_number_from_inventory()

    def get_node_number_from_inventory(self):
        node_number_list = []
        for line in self.inventory_lines:
            match = re.search(self.node_pattern, line)
            if match:
                node_number_list.append(int(match.group(0)[5:-1]))
        #sorted_node_number_list = sorted(node_number_list)
        return node_number_list[-1] + 1    # node 从0开始计数

    def generate_inventory(self, host_list: list, private_key: str, nginx_port: int):
        if self.node_number > len(host_list):
            print(f"host节点数不足，至少需要{self.node_number}个，请重新输入")
            raise ValueError("host_list not enough")
        # host列表排序
        # host_list = sorted(host_list)
        with open(self.output_inventory_file, 'w', encoding='utf-8') as output:
            line_index = 1
            for line in self.inventory_lines:
                line_index = line_index + 1
                if "{ssh_private_key}" in line:
                    line = line.replace("{ssh_private_key}", private_key)
                    output.write(line)
                    continue
                if "{nginx_port}" in line:
                    line = line.replace("{nginx_port}", str(nginx_port))
                    output.write(line)
                    continue
                match = re.search(self.node_pattern, line)
                try:
                    if match:
                        node_name = match.group(0)
                        node_index = int(node_name[5:-1])    # {node(index)}去掉前后的符号
                        line = line.replace(node_name, host_list[node_index])
                        output.write(line)
                        continue
                except Exception as e:
                    raise ValueError(f"line {line_index}: replace node_host failed, error: {str(e)}, line: {line}")
                output.write(line)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--host_list', required=False, type=str, default="192.168.0.94,192.168.0.120,192.168.0.73,192.168.0.118,192.168.0.58,192.168.0.212,192.168.0.223,192.168.0.182")
    parser.add_argument('--inventory_template_file', required=False, type=str, default="omni_infer_inventory_template.yml")
    parser.add_argument('--output_inventory_file', required=False, type=str, default="omni_infer_inventory.yml")
    parser.add_argument('--nginx_port', required=False, type=int, default=7000)
    parser.add_argument('--node_number', required=False, type=int, default=4)
    parser.add_argument('--private_key', required=False, type=str, default="/data/CI/pem/key.pem")
    args = parser.parse_args()

    host_list = args.host_list.split(',')
    inventory = Inventory(args.inventory_template_file, args.output_inventory_file)
    inventory.generate_inventory(host_list, args.private_key, args.nginx_port)

